{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Categorical Regression\n",
    "\n",
    "In this example, we will use the categorical family to model outcomes with more than two categories. The examples in this notebook were constructed by Tomás Capretto, and assembled into this example by Tyler James Burch ([\\@tjburch](https://github.com/tjburch) on GitHub)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "import arviz as az\n",
    "import bambi as bmb\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "\n",
    "warnings.simplefilter(action=\"ignore\", category=FutureWarning)\n",
    "\n",
    "az.style.use(\"arviz-darkgrid\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When modeling binary outcomes with Bambi, the Bernoulli family is used. The multivariate generalization of the Bernoulli family is the Categorical family, and with it, we can model an arbitrary number of outcome categories. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example with toy data\n",
    "\n",
    "To start, we will create a toy dataset with three classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(1234)\n",
    "x = np.hstack([rng.normal(m, s, size=50) for m, s in zip([-2.5, 0, 2.5], [1.2, 0.5, 1.2])])\n",
    "y = np.array([\"A\"] * 50 + [\"B\"] * 50 + [\"C\"] * 50)\n",
    "\n",
    "colors = [\"C0\"] * 50 + [\"C1\"] * 50 + [\"C2\"] * 50\n",
    "plt.scatter(x, np.random.uniform(size=150), color=colors)\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\");"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we have 3 classes, generated from three normal distributions: $N(-2.5, 1.2)$, $N(0, 0.5)$, and $N(2.5, 1.2)$. Creating a model to fit these distributions,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.DataFrame({\"y\": y, \"x\": x})\n",
    "model = bmb.Model(\"y ~ x\", data, family=\"categorical\")\n",
    "idata = model.fit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we pass the `family=\"categorical\"` argument to Bambi's `Model` method in order to call the categorical family. Here, the response variable are strings (\"A\", \"B\", \"C\"), however they can also be `pd.Categorical` objects.\n",
    "\n",
    "Next we will use posterior predictions to visualize the mean class probability across the $x$ spectrum.\n",
    "To do this, we simply call [bmb.interpret.plot_predictions](https://bambinos.github.io/bambi/notebooks/plot_predictions.html) and specify that the computations should be conditional on $x$. Then, the function does all the work for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = bmb.interpret.plot_predictions(model, idata, conditional=\"x\")\n",
    "ax[0].set(ylabel=\"Probability of category\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want a more elaborate figure, we can tweak the function arguments like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = bmb.interpret.plot_predictions(\n",
    "    model,\n",
    "    idata,\n",
    "    conditional=\"x\",\n",
    "    legend=False,\n",
    "    subplot_kwargs={\"main\": \"x\", \"group\": \"estimate_dim\", \"panel\": \"estimate_dim\"},\n",
    "    fig_kwargs={\"figsize\": (10, 5), \"sharey\": True}\n",
    ")\n",
    "ax[0].set(ylabel=\"Probability\")\n",
    "ax[1].set(ylabel=None)\n",
    "ax[2].set(ylabel=None);"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we can notice that the probability phases between classes from left to right. At all points across $x$, sum of the class probabilities is 1, since in our generative model, it must be one of these three outcomes.  "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The iris dataset\n",
    "\n",
    "Next, we will look at the classic \"iris\" dataset, which contains samples from 3 different species of iris plants. Using properties of the plant, we will try to model its species."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_iris = sns.load_dataset(\"iris\")\n",
    "df_iris.head(3)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset includes four different properties of the plants: it's sepal length, sepal width, petal length, and petal width. There are 3 different class possibilities: setosa, versicolor, and virginica."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.pairplot(df_iris, hue=\"species\");"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the three species have several distinct characteristics, which our linear model can capture to distinguish between them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = bmb.Model(\n",
    "    \"species ~ sepal_length + sepal_width + petal_length + petal_width\",\n",
    "    df_iris,\n",
    "    family=\"categorical\",\n",
    ")\n",
    "idata = model.fit()\n",
    "az.summary(idata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_trace(idata);"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that this has fit quite nicely. You'll notice there are $n-1$ parameters to fit, where $n$ is the number of categories. In the minimal binary case, recall there's only one parameter set, since it models probability $p$ of being in a class, and probability $1-p$ of being in the other class. Using the categorical distribution, this extends, so we have $p_1$ for class 1, $p_2$ for class 2, and $1-(p_1+p_2)$ for the final class."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using numerical and categorical predictors\n",
    "\n",
    "Next we will look at an example from chapter 8 of Alan Agresti's _Categorical Data Analysis_, looking at the primary food choice for 64 alligators caught in Lake George, Florida. We will use their length (a continuous variable) and sex (a categorical variable) as predictors to model their food choice.\n",
    "\n",
    "First, reproducing the dataset,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "length = [\n",
    "    1.3, 1.32, 1.32, 1.4, 1.42, 1.42, 1.47, 1.47, 1.5, 1.52, 1.63, 1.65, 1.65, 1.65, 1.65,\n",
    "    1.68, 1.7, 1.73, 1.78, 1.78, 1.8, 1.85, 1.93, 1.93, 1.98, 2.03, 2.03, 2.31, 2.36, 2.46,\n",
    "    3.25, 3.28, 3.33, 3.56, 3.58, 3.66, 3.68, 3.71, 3.89, 1.24, 1.3, 1.45, 1.45, 1.55, 1.6,\n",
    "    1.6, 1.65, 1.78, 1.78, 1.8, 1.88, 2.16, 2.26, 2.31, 2.36, 2.39, 2.41, 2.44, 2.56, 2.67,\n",
    "    2.72, 2.79, 2.84\n",
    "]\n",
    "choice = [\n",
    "    \"I\", \"F\", \"F\", \"F\", \"I\", \"F\", \"I\", \"F\", \"I\", \"I\", \"I\", \"O\", \"O\", \"I\", \"F\", \"F\",\n",
    "    \"I\", \"O\", \"F\", \"O\", \"F\", \"F\", \"I\", \"F\", \"I\", \"F\", \"F\", \"F\", \"F\", \"F\", \"O\", \"O\",\n",
    "    \"F\", \"F\", \"F\", \"F\", \"O\", \"F\", \"F\", \"I\", \"I\", \"I\", \"O\", \"I\", \"I\", \"I\", \"F\", \"I\",\n",
    "    \"O\", \"I\", \"I\", \"F\", \"F\", \"F\", \"F\", \"F\", \"F\", \"F\", \"O\", \"F\", \"I\", \"F\", \"F\"\n",
    "]\n",
    "\n",
    "sex = [\"Male\"] * 32 + [\"Female\"] * 31\n",
    "data = pd.DataFrame({\"choice\": choice, \"length\": length, \"sex\": sex})\n",
    "data[\"choice\"]  = pd.Categorical(\n",
    "    data[\"choice\"].map({\"I\": \"Invertebrates\", \"F\": \"Fish\", \"O\": \"Other\"}),\n",
    "    [\"Other\", \"Invertebrates\", \"Fish\"],\n",
    "    ordered=True\n",
    ")\n",
    "data.head(3)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, constructing the model,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = bmb.Model(\"choice ~ length + sex\", data, family=\"categorical\")\n",
    "idata = model.fit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using [bmb.interpret.plot_predictions](https://bambinos.github.io/bambi/notebooks/plot_predictions.html), we can visualize how the probability of the different response levels varies conditional on a set of predictors. In the plot below, we visualize how the food choices vary by length for both male and female alligators. Note how `estimate_dim` (the response level) is mapped as the value to the `group` key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bmb.interpret.plot_predictions(\n",
    "    model,\n",
    "    idata,\n",
    "    [\"length\", \"sex\"],\n",
    "    subplot_kwargs={\"main\": \"length\", \"group\": \"estimate_dim\", \"panel\": \"sex\"},\n",
    "    fig_kwargs={\"figsize\": (12, 4)},\n",
    "    legend=True\n",
    ");"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we can see that the larger male and female alligators are, the less of a taste they have for invertebrates, and far prefer fish. Additionally, males seem to have a higher propensity to consume \"other\" foods compared to females at any size. Of note, the posterior means predicted by Bambi contain information about all $n$ categories (despite having only $n-1$ coefficients), so we can directly construct this plot, rather than manually calculating $1-(p_1+p_2)$ for the third class.\n",
    "\n",
    "Last, we can make a posterior predictive plot,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.predict(idata, kind=\"pps\")\n",
    "\n",
    "ax = az.plot_ppc(idata)\n",
    "ax.set_xticks([0.5, 1.5, 2.5])\n",
    "ax.set_xticklabels(model.response_component.term.levels)\n",
    "ax.set_xlabel(\"Choice\");\n",
    "ax.set_ylabel(\"Probability\");"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which depicts posterior predicted probability for each possible food choice for an alligator, which reinforces fish being the most likely food choice, followed by invertebrates."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "Agresti, A. (2013) Categorical Data Analysis. 3rd Edition, John Wiley & Sons Inc., Hoboken."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext watermark\n",
    "%watermark -n -u -v -iv -w"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dev",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
